use num_bigint::BigUint;
use num_bigint::RandBigInt;
use num_bigint::ToBigUint;
use num_primes::Generator;
use num_traits::One;

fn generate_safe_prime(prime_bits: usize) -> BigUint {
    let c = Generator::safe_prime(prime_bits); // use num-primes crate
    let d = c.clone().to_bytes_be();
    let f = BigUint::from_bytes_be(&d); // use num-bigint crate
    f
}

fn main() {
    // p and g are public; one can usually request them from server
    let p = generate_safe_prime(512);
    let g = 2.to_biguint().unwrap();

    let mut rng = rand::thread_rng();
    let secret_a = rng.gen_biguint(2048); // A's secret
    let secret_b = rng.gen_biguint(2048); // B's secret

    let g_a = g.modpow(&secret_a, &p); // A sends g_a to B via a public message
    let g_b = g.modpow(&secret_b, &p); // B sends g_b to A via a public message

    let shared_secret_computed_by_A = g_b.modpow(&secret_a, &p);
    let shared_secret_computed_by_B = g_a.modpow(&secret_b, &p);

    println!(
        "shared secret computed by A: {}",
        &shared_secret_computed_by_A
    );
    println!(
        "shared secret computed by B: {}",
        &shared_secret_computed_by_B
    );
    assert_eq!(&shared_secret_computed_by_A, &shared_secret_computed_by_B);

    /*Output
    shared secret computed by A: 7085017045441200181467502439139014522983597461003612144107748514118266429715983022633256410
    640002617261361476402978219402626623314766855257317520461001895
    shared secret computed by B: 7085017045441200181467502439139014522983597461003612144107748514118266429715983022633256410
    640002617261361476402978219402626623314766855257317520461001895
    */
}
